

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <erl_nif.h>
#include "opus.h"

#define FRAME_SIZE 960 // 20ms 的音频数据
#define SAMPLE_RATE 16000
#define CHANNELS 1
#define MAX_PACKET_SIZE 1276 // Opus 编码后的最大数据大小

//  gcc -fPIC -shared  -o priv/opus_nif.so priv/opus_nif.c -I/usr/include/opus -lopus
// 定义结构体保存编码器和解码器
typedef struct
{
    OpusEncoder *encoder;
    OpusDecoder *decoder;
    // unsigned char *pcm_buffer; // 缓存 PCM 数据
    // size_t pcm_buffer_size;    // 当前缓存的大小
} CodecState;

// NIF 加载函数
static int load(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM load_info)
{
    int error;
    CodecState *state = (CodecState *)enif_alloc(sizeof(CodecState));
    if (!state)
    {
        fprintf(stderr, "memory allocation error\n");
        return -1;
    }

    // 初始化 Opus 编码器
    state->encoder = opus_encoder_create(SAMPLE_RATE, CHANNELS, OPUS_APPLICATION_AUDIO, &error);
    if (error != OPUS_OK)
    {
        fprintf(stderr, "encoder initialization error: %d\n", error);
        enif_free(state);
        return -1;
    }

    // 初始化 Opus 解码器
    state->decoder = opus_decoder_create(SAMPLE_RATE, CHANNELS, &error);
    if (error != OPUS_OK)
    {
        fprintf(stderr, "decoder initialization error: %d\n", error);
        opus_encoder_destroy(state->encoder);
        enif_free(state);
        return -1;
    }
    
    // 初始化缓存
    // state->pcm_buffer = (unsigned char *)enif_alloc(FRAME_SIZE * CHANNELS * sizeof(short));
    // if (!state->pcm_buffer)
    // {
    //     fprintf(stderr, "memory allocation error for PCM buffer\n");
    //     opus_encoder_destroy(state->encoder);
    //     opus_decoder_destroy(state->decoder);
    //     enif_free(state);
    //     return -1;
    // }
    // state->pcm_buffer_size = 0;

    *priv_data = state;
    fprintf(stderr, "NIF initialized successfully\n");
    return 0;
}

// NIF 卸载函数
static void unload(ErlNifEnv *env, void *priv_data)
{
    CodecState *state = (CodecState *)priv_data;
    if (state)
    {
        if (state->encoder)
        {
            opus_encoder_destroy(state->encoder);
        }
        if (state->decoder)
        {
            opus_decoder_destroy(state->decoder);
        }
        // if (state->pcm_buffer) {
        //     enif_free(state->pcm_buffer);
        // }
        enif_free(state);
    }
}

// NIF 函数：编码 PCM 为 Opus 数据
static ERL_NIF_TERM encode_pcm_data(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    if (argc != 1)
    {
        return enif_make_badarg(env);
    }

    ErlNifBinary pcm_bin;
    if (!enif_inspect_binary(env, argv[0], &pcm_bin))
    {
        return enif_make_badarg(env);
    }

    if (pcm_bin.size % (FRAME_SIZE * CHANNELS * sizeof(short)) != 0)
    {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "invalid_pcm_size"));
    }

    CodecState *state = (CodecState *)enif_priv_data(env);
    if (!state || !state->encoder)
    {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "encoder_not_initialized"));
    }

    // if (pcm_bin.size % (CHANNELS * sizeof(short)) != 0) {
    //     return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "invalid_pcm_size"));
    // }

    short *pcm = (opus_int16 *)pcm_bin.data;
    int frame_count = pcm_bin.size / (CHANNELS * sizeof(short));

    // 检查 PCM 数据大小是否是 FRAME_SIZE 的整数倍

    if (frame_count % FRAME_SIZE != 0)
    {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "invalid_pcm_size"));
    }

    unsigned char opus_data[MAX_PACKET_SIZE];
    int encoded_size = opus_encode(state->encoder, pcm, frame_count, opus_data, MAX_PACKET_SIZE);
    // if (encoded_size < 0) {
    //     fprintf(stderr, "encode error: %d\n", encoded_size);
    //     return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, encoded_size));
    // }

    if (encoded_size < 0)
    {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_string(env, opus_strerror(encoded_size), ERL_NIF_LATIN1));
    }

    // ErlNifBinary result;
    // enif_alloc_binary(encoded_size, &result);
    // memcpy(result.data, opus_data, encoded_size);
    // return enif_make_binary(env, &result);
    ERL_NIF_TERM opus_bin;
    unsigned char *encoded_data = enif_make_new_binary(env, encoded_size, &opus_bin);
    memcpy(encoded_data, opus_data, encoded_size);

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), opus_bin);
}

// static ERL_NIF_TERM encode_pcm_data(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
//     if (argc != 1) {
//         return enif_make_badarg(env);
//     }

//     ErlNifBinary pcm_bin;
//     if (!enif_inspect_binary(env, argv[0], &pcm_bin)) {
//         return enif_make_badarg(env);
//     }

//     CodecState *state = (CodecState *)enif_priv_data(env);
//     if (!state || !state->encoder) {
//         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "encoder_not_initialized"));
//     }

//     // 将新数据追加到缓存
//     size_t new_data_size = pcm_bin.size;
//     unsigned char *new_data = pcm_bin.data;

//     // 计算剩余空间
//     size_t remaining_space = (FRAME_SIZE * CHANNELS * sizeof(short)) - state->pcm_buffer_size;

//     // 如果新数据可以完全放入缓存
//     if (new_data_size <= remaining_space) {
//         memcpy(state->pcm_buffer + state->pcm_buffer_size, new_data, new_data_size);
//         state->pcm_buffer_size += new_data_size;

//         // 如果缓存未满，返回等待更多数据
//         if (state->pcm_buffer_size < FRAME_SIZE * CHANNELS * sizeof(short)) {
//             return enif_make_tuple2(env, enif_make_atom(env, "ok"), enif_make_atom(env, "wait_for_more_data"));
//         }
//     } else {
//         // 如果新数据超出缓存空间，先处理缓存中的数据
//         memcpy(state->pcm_buffer + state->pcm_buffer_size, new_data, remaining_space);
//         state->pcm_buffer_size += remaining_space;

//         // 剩余数据需要等待下一轮
//         new_data += remaining_space;
//         new_data_size -= remaining_space;
//     }

//     // 编码缓存中的数据
//     short *pcm = (opus_int16 *)state->pcm_buffer;
//     unsigned char opus_data[MAX_PACKET_SIZE];
//     int encoded_size = opus_encode(state->encoder, pcm, FRAME_SIZE, opus_data, MAX_PACKET_SIZE);

//     if (encoded_size < 0) {
//         return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_string(env, opus_strerror(encoded_size), ERL_NIF_LATIN1));
//     }

//     // 返回编码后的数据
//     ERL_NIF_TERM opus_bin;
//     unsigned char *encoded_data = enif_make_new_binary(env, encoded_size, &opus_bin);
//     memcpy(encoded_data, opus_data, encoded_size);

//     // 重置缓存
//     state->pcm_buffer_size = 0;

//     return enif_make_tuple2(env, enif_make_atom(env, "ok"), opus_bin);
// }

// NIF 函数：解码 Opus 数据
static ERL_NIF_TERM decode_opus_data(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
{
    if (argc != 1)
    {
        return enif_make_badarg(env);
    }

    ErlNifBinary opus_bin;
    if (!enif_inspect_binary(env, argv[0], &opus_bin))
    {
        return enif_make_badarg(env);
    }

    CodecState *state = (CodecState *)enif_priv_data(env);
    if (!state || !state->decoder)
    {
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_atom(env, "decoder_not_initialized"));
    }

    short pcm[FRAME_SIZE * CHANNELS];
    int samples_decoded = opus_decode(state->decoder, opus_bin.data, opus_bin.size, pcm, FRAME_SIZE, 0);
    if (samples_decoded < 0)
    {
        fprintf(stderr, "解码失败: %d\n", samples_decoded);
        return enif_make_tuple2(env, enif_make_atom(env, "error"), enif_make_int(env, samples_decoded));
    }

    ERL_NIF_TERM pcm_bin;
    unsigned char *pcm_data = enif_make_new_binary(env, samples_decoded * CHANNELS * sizeof(short), &pcm_bin);
    memcpy(pcm_data, pcm, samples_decoded * CHANNELS * sizeof(short));

    return enif_make_tuple2(env, enif_make_atom(env, "ok"), pcm_bin);
}


// NIF 函数表
static ErlNifFunc nif_funcs[] = {
    {"encode_pcm_data", 1, encode_pcm_data},
    {"decode_opus_data", 1, decode_opus_data}};

// NIF 模块初始化
ERL_NIF_INIT(Elixir.SockWorkstation.CodecNifGen, nif_funcs, &load, NULL, NULL, &unload)
